{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50e1d1d5-3bdd-4224-9f93-bf5d9a83f424",
   "metadata": {},
   "source": [
    "# 使用 OpenVINO 创建由 LLM 驱动的聊天机器人\n",
    "\n",
    "在快速发展的人工智能 (AI) 世界中，聊天机器人已成为企业增强客户互动和简化运营的强大工具。 大型语言模型 (LLM) 是能够理解和生成人类语言的人工智能系统。 他们使用深度学习算法和大量数据来学习语言的细微差别并产生连贯且相关的响应。 虽然一个不错的基于意图的聊天机器人可以回答基本的一键式查询，例如订单管理、常见问题解答和政策问题，但 LLM 聊天机器人可以解决更复杂的多点触控问题。 LLM 使聊天机器人能够通过上下文记忆以对话方式提供支持，类似于人类的做法。 利用语言模型的功能，聊天机器人变得越来越智能，能够以极高的准确性理解和响应人类语言。\n",
    "\n",
    "在本教程中，我们将考虑如何使用 OpenVINO 的强大功能来运行大型语言模型以进行聊天。 我们将使用 ModelScope 库中的预训练模型。 为了简化用户体验，使用 Hugging Face Optimum Intel 库将模型转换为 OpenVINO™ IR 格式。\n",
    "\n",
    "\n",
    "\n",
    "本教程包含以下步骤:\n",
    "\n",
    "- 安装先决条件\n",
    "- 使用 OpenVINO 与 Hugging Face Optimum 集成从公共源下载并转换模型 [OpenVINO integration with Hugging Face Optimum](https://huggingface.co/blog/openvino).\n",
    "- 使用 NNCF 将模型权重压缩为 4 位或 8 位数据类型 [NNCF](https://github.com/openvinotoolkit/nncf)\n",
    "- 创建聊天推理管道\n",
    "- 运行聊天pipeline\n",
    "\n",
    "\n",
    "#### 目录:\n",
    "\n",
    "- [先决条件](#Prerequisites)\n",
    "- [选择模型进行推理](#Select-model-for-inference)\n",
    "- [访问modelscope上的模型](#login-to-huggingfacehub-to-get-access-to-pretrained-model)\n",
    "- [使用 Optimum Intel 实例化模型](#Instantiate-Model-using-Optimum-Intel)\n",
    "- [压缩模型权重](#Compress-model-weights)\n",
    "    - [使用 Optimum Intel 进行权重压缩](#Weights-Compression-using-Optimum-Intel)\n",
    "    - [使用 NNCF 进行权重压缩](#Weights-Compression-using-NNCF)\n",
    "- [Select device for inference and model variant](#Select-device-for-inference-and-model-variant)\n",
    "- [Run Chatbot](#Run-Chatbot)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5df233b0-0369-4fff-9952-7957a90394a5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Prerequisites\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Install required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "563ecf9f-346b-4f14-85ef-c66ff0c95f65",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-03-14T03:05:41.650053Z",
     "iopub.status.busy": "2024-03-14T03:05:41.649428Z",
     "iopub.status.idle": "2024-03-14T03:09:36.825005Z",
     "shell.execute_reply": "2024-03-14T03:09:36.823554Z",
     "shell.execute_reply.started": "2024-03-14T03:05:41.649995Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "ms-swift 1.5.4 requires transformers<4.38,>=4.33, but you have transformers 4.38.2 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu\\\n",
    "\"git+https://github.com/huggingface/optimum-intel.git\"\\\n",
    "\"git+https://github.com/openvinotoolkit/nncf.git\"\\\n",
    "\"datasets\" \\\n",
    "\"accelerate\"\\\n",
    "\"openvino-nightly\"\\\n",
    "\"gradio\"\\\n",
    "\"onnx\" \"einops\" \"transformers_stream_generator\" \"tiktoken\" \"transformers>=4.38.1\" \"bitsandbytes\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81983176-e571-4652-ba21-4bd608c35146",
   "metadata": {},
   "source": [
    "## Select model for inference\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The tutorial supports different models, you can select one from the provided options to compare the quality of open source LLM solutions.\n",
    ">**Note**: conversion of some models can require additional actions from user side and at least 64GB RAM for conversion.\n",
    "\n",
    "The available options are:\n",
    "\n",
    "* **tiny-llama-1b-chat** - This is the chat model finetuned on top of [TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T](https://huggingface.co/TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T). The TinyLlama project aims to pretrain a 1.1B Llama model on 3 trillion tokens with the adoption of the same architecture and tokenizer as Llama 2. This means TinyLlama can be plugged and played in many open-source projects built upon Llama. Besides, TinyLlama is compact with only 1.1B parameters. This compactness allows it to cater to a multitude of applications demanding a restricted computation and memory footprint. More details about model can be found in [model card](https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0)\n",
    "*  **mini-cpm-2b-dpo** - MiniCPM is an End-Size LLM developed by ModelBest Inc. and TsinghuaNLP, with only 2.4B parameters excluding embeddings. After Direct Preference Optimization (DPO) fine-tuning, MiniCPM outperforms many popular 7b, 13b and 70b models. More details can be found in [model_card](https://huggingface.co/openbmb/MiniCPM-2B-dpo-fp16).\n",
    "*  **gemma-2b-it** - Gemma is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. They are text-to-text, decoder-only large language models, available in English, with open weights, pre-trained variants, and instruction-tuned variants. Gemma models are well-suited for a variety of text generation tasks, including question answering, summarization, and reasoning. This model is instruction-tuned version of 2B parameters model. More details about model can be found in [model card](https://huggingface.co/google/gemma-2b-it).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/google/gemma-2b-it), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "* **red-pajama-3b-chat** - A 2.8B parameter pre-trained language model based on GPT-NEOX architecture. It was developed by Together Computer and leaders from the open-source AI community. The model is fine-tuned on OASST1 and Dolly2 datasets to enhance chatting ability. More details about model can be found in [HuggingFace model card](https://huggingface.co/togethercomputer/RedPajama-INCITE-Chat-3B-v1).\n",
    "*  **gemma-7b-it** - Gemma is a family of lightweight, state-of-the-art open models from Google, built from the same research and technology used to create the Gemini models. They are text-to-text, decoder-only large language models, available in English, with open weights, pre-trained variants, and instruction-tuned variants. Gemma models are well-suited for a variety of text generation tasks, including question answering, summarization, and reasoning. This model is instruction-tuned version of 7B parameters model. More details about model can be found in [model card](https://huggingface.co/google/gemma-7b-it).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/google/gemma-7b-it), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "\n",
    "* **llama-2-7b-chat** - LLama 2 is the second generation of LLama models developed by Meta. Llama 2 is a collection of pre-trained and fine-tuned generative text models ranging in scale from 7 billion to 70 billion parameters. llama-2-7b-chat is 7 billions parameters version of LLama 2 finetuned and optimized for dialogue use case. More details about model can be found in the [paper](https://ai.meta.com/research/publications/llama-2-open-foundation-and-fine-tuned-chat-models/), [repository](https://github.com/facebookresearch/llama) and [HuggingFace model card](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf).\n",
    ">**Note**: run model with demo, you will need to accept license agreement. \n",
    ">You must be a registered user in 🤗 Hugging Face Hub. Please visit [HuggingFace model card](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf), carefully read terms of usage and click accept button.  You will need to use an access token for the code below to run. For more information on access tokens, refer to [this section of the documentation](https://huggingface.co/docs/hub/security-tokens).\n",
    ">You can login on Hugging Face Hub in notebook environment, using following code:\n",
    " \n",
    "```python\n",
    "    ## login to huggingfacehub to get access to pretrained model \n",
    "\n",
    "    from huggingface_hub import notebook_login, whoami\n",
    "\n",
    "    try:\n",
    "        whoami()\n",
    "        print('Authorization token already provided')\n",
    "    except OSError:\n",
    "        notebook_login()\n",
    "```\n",
    "* **qwen1.5-0.5b-chat/qwen1.5-1.8b-chat/qwen1.5-7b-chat** - Qwen1.5 is the beta version of Qwen2, a transformer-based decoder-only language model pretrained on a large amount of data. Qwen1.5 is a language model series including decoder language models of different model sizes. It is based on the Transformer architecture with SwiGLU activation, attention QKV bias, group query attention, mixture of sliding window attention and full attention. You can find more details about model in the [model repository](https://huggingface.co/Qwen).\n",
    "* **qwen-7b-chat** - Qwen-7B is the 7B-parameter version of the large language model series, Qwen (abbr. Tongyi Qianwen), proposed by Alibaba Cloud. Qwen-7B is a Transformer-based large language model, which is pretrained on a large volume of data, including web texts, books, codes, etc. For more details about Qwen, please refer to the [GitHub](https://github.com/QwenLM/Qwen) code repository.\n",
    "* **mpt-7b-chat** - MPT-7B is part of the family of MosaicPretrainedTransformer (MPT) models, which use a modified transformer architecture optimized for efficient training and inference. These architectural changes include performance-optimized layer implementations and the elimination of context length limits by replacing positional embeddings with Attention with Linear Biases ([ALiBi](https://arxiv.org/abs/2108.12409)). Thanks to these modifications, MPT models can be trained with high throughput efficiency and stable convergence. MPT-7B-chat is a chatbot-like model for dialogue generation. It was built by finetuning MPT-7B on the [ShareGPT-Vicuna](https://huggingface.co/datasets/jeffwan/sharegpt_vicuna), [HC3](https://huggingface.co/datasets/Hello-SimpleAI/HC3), [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca), [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf), and [Evol-Instruct](https://huggingface.co/datasets/victor123/evol_instruct_70k) datasets. More details about the model can be found in [blog post](https://www.mosaicml.com/blog/mpt-7b), [repository](https://github.com/mosaicml/llm-foundry/) and [HuggingFace model card](https://huggingface.co/mosaicml/mpt-7b-chat).\n",
    "* **chatglm3-6b** - ChatGLM3-6B is the latest open-source model in the ChatGLM series. While retaining many excellent features such as smooth dialogue and low deployment threshold from the previous two generations, ChatGLM3-6B employs a more diverse training dataset, more sufficient training steps, and a more reasonable training strategy. ChatGLM3-6B adopts a newly designed [Prompt format](https://github.com/THUDM/ChatGLM3/blob/main/PROMPT_en.md), in addition to the normal multi-turn dialogue. You can find more details about model in the [model card](https://huggingface.co/THUDM/chatglm3-6b)\n",
    "* **mistral-7b** - The Mistral-7B-v0.1 Large Language Model (LLM) is a pretrained generative text model with 7 billion parameters. You can find more details about model in the [model card](https://huggingface.co/mistralai/Mistral-7B-v0.1), [paper](https://arxiv.org/abs/2310.06825) and [release blog post](https://mistral.ai/news/announcing-mistral-7b/).\n",
    "* **zephyr-7b-beta** - Zephyr is a series of language models that are trained to act as helpful assistants. Zephyr-7B-beta is the second model in the series, and is a fine-tuned version of [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1) that was trained on on a mix of publicly available, synthetic datasets using [Direct Preference Optimization (DPO)](https://arxiv.org/abs/2305.18290). You can find more details about model in [technical report](https://arxiv.org/abs/2310.16944) and [HuggingFace model card](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta).\n",
    "* **neural-chat-7b-v3-1** - Mistral-7b model fine-tuned using Intel Gaudi. The model fine-tuned on the open source dataset [Open-Orca/SlimOrca](https://huggingface.co/datasets/Open-Orca/SlimOrca) and aligned with [Direct Preference Optimization (DPO) algorithm](https://arxiv.org/abs/2305.18290). More details can be found in [model card](https://huggingface.co/Intel/neural-chat-7b-v3-1) and [blog post](https://medium.com/@NeuralCompressor/the-practice-of-supervised-finetuning-and-direct-preference-optimization-on-habana-gaudi2-a1197d8a3cd3).\n",
    "* **notus-7b-v1** - Notus is a collection of fine-tuned models using [Direct Preference Optimization (DPO)](https://arxiv.org/abs/2305.18290). and related [RLHF](https://huggingface.co/blog/rlhf) techniques. This model is the first version, fine-tuned with DPO over zephyr-7b-sft. Following a data-first approach, the only difference between Notus-7B-v1 and Zephyr-7B-beta is the preference dataset used for dDPO. Proposed approach for dataset creation helps to effectively fine-tune Notus-7b that surpasses Zephyr-7B-beta and Claude 2 on [AlpacaEval](https://tatsu-lab.github.io/alpaca_eval/). More details about model can be found in [model card](https://huggingface.co/argilla/notus-7b-v1).\n",
    "* **youri-7b-chat** - Youri-7b-chat is a Llama2 based model. [Rinna Co., Ltd.](https://rinna.co.jp/) conducted further pre-training for the Llama2 model with a mixture of English and Japanese datasets to improve Japanese task capability. The model is publicly released on Hugging Face hub. You can find detailed information at the [rinna/youri-7b-chat project page](https://huggingface.co/rinna/youri-7b). \n",
    "* **baichuan2-7b-chat** - Baichuan 2 is the new generation of large-scale open-source language models launched by [Baichuan Intelligence inc](https://www.baichuan-ai.com/home). It is trained on a high-quality corpus with 2.6 trillion tokens and has achieved the best performance in authoritative Chinese and English benchmarks of the same size.\n",
    "* **internlm2-chat-1.8b** - InternLM2 is the second generation InternLM series. Compared to the previous generation model, it shows significant improvements in various capabilities, including reasoning, mathematics, and coding. More details about model can be found in [model repository](https://huggingface.co/internlm).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f93282b6-f1f1-4153-84af-31aac79c3ef4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:26:56.297204Z",
     "iopub.status.busy": "2024-03-14T03:26:56.296646Z",
     "iopub.status.idle": "2024-03-14T03:27:00.889296Z",
     "shell.execute_reply": "2024-03-14T03:27:00.888046Z",
     "shell.execute_reply.started": "2024-03-14T03:26:56.297159Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: ipywidgets in /opt/conda/lib/python3.10/site-packages (8.1.2)\n",
      "Requirement already satisfied: comm>=0.1.3 in /opt/conda/lib/python3.10/site-packages (from ipywidgets) (0.2.1)\n",
      "Requirement already satisfied: ipython>=6.1.0 in /opt/conda/lib/python3.10/site-packages (from ipywidgets) (8.19.0)\n",
      "Requirement already satisfied: traitlets>=4.3.1 in /opt/conda/lib/python3.10/site-packages (from ipywidgets) (5.14.1)\n",
      "Requirement already satisfied: widgetsnbextension~=4.0.10 in /opt/conda/lib/python3.10/site-packages (from ipywidgets) (4.0.10)\n",
      "Requirement already satisfied: jupyterlab-widgets~=3.0.10 in /opt/conda/lib/python3.10/site-packages (from ipywidgets) (3.0.10)\n",
      "Requirement already satisfied: decorator in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (4.4.2)\n",
      "Requirement already satisfied: jedi>=0.16 in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.19.1)\n",
      "Requirement already satisfied: matplotlib-inline in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.1.6)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>=3.0.41 in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (3.0.43)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (2.17.2)\n",
      "Requirement already satisfied: stack-data in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (0.6.3)\n",
      "Requirement already satisfied: exceptiongroup in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (1.2.0)\n",
      "Requirement already satisfied: pexpect>4.3 in /opt/conda/lib/python3.10/site-packages (from ipython>=6.1.0->ipywidgets) (4.9.0)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.3 in /opt/conda/lib/python3.10/site-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets) (0.8.3)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /opt/conda/lib/python3.10/site-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /opt/conda/lib/python3.10/site-packages (from prompt-toolkit<3.1.0,>=3.0.41->ipython>=6.1.0->ipywidgets) (0.2.12)\n",
      "Requirement already satisfied: executing>=1.2.0 in /opt/conda/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.0.1)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /opt/conda/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.4.1)\n",
      "Requirement already satisfied: pure-eval in /opt/conda/lib/python3.10/site-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.2.2)\n",
      "Requirement already satisfied: six>=1.12.0 in /opt/conda/lib/python3.10/site-packages (from asttokens>=2.1.0->stack-data->ipython>=6.1.0->ipywidgets) (1.16.0)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install ipywidgets\n",
    "from config import SUPPORTED_LLM_MODELS\n",
    "import ipywidgets as widgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e02b34fb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:27:06.767156Z",
     "iopub.status.busy": "2024-03-14T03:27:06.766479Z",
     "iopub.status.idle": "2024-03-14T03:27:06.780548Z",
     "shell.execute_reply": "2024-03-14T03:27:06.779495Z",
     "shell.execute_reply.started": "2024-03-14T03:27:06.767098Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a4d0977b6c345cebc841d91773ed51f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model Language:', options=('English', 'Chinese', 'Japanese'), value='English')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_languages = list(SUPPORTED_LLM_MODELS)\n",
    "\n",
    "model_language = widgets.Dropdown(\n",
    "    options=model_languages,\n",
    "    value=model_languages[0],\n",
    "    description=\"Model Language:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8d22fedb-d1f6-4306-b910-efac5b849c7c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:27:12.285453Z",
     "iopub.status.busy": "2024-03-14T03:27:12.284872Z",
     "iopub.status.idle": "2024-03-14T03:27:12.295707Z",
     "shell.execute_reply": "2024-03-14T03:27:12.294538Z",
     "shell.execute_reply.started": "2024-03-14T03:27:12.285411Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc8e48b7e59242df946b4dda973b4b2e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model:', options=('qwen1.5-0.5b-chat', 'qwen1.5-1.8b-chat', 'qwen1.5-7b-chat', 'qwen-7b-…"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_ids = list(SUPPORTED_LLM_MODELS[model_language.value])\n",
    "\n",
    "model_id = widgets.Dropdown(\n",
    "    options=model_ids,\n",
    "    value=model_ids[0],\n",
    "    description=\"Model:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "906022ec-96bf-41a9-9447-789d2e875250",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:27:16.263275Z",
     "iopub.status.busy": "2024-03-14T03:27:16.262685Z",
     "iopub.status.idle": "2024-03-14T03:27:16.269078Z",
     "shell.execute_reply": "2024-03-14T03:27:16.267846Z",
     "shell.execute_reply.started": "2024-03-14T03:27:16.263235Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected model qwen1.5-0.5b-chat\n"
     ]
    }
   ],
   "source": [
    "model_configuration = SUPPORTED_LLM_MODELS[model_language.value][model_id.value]\n",
    "print(f\"Selected model {model_id.value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62af3e8a-915a-49b4-8007-803777ba9eaf",
   "metadata": {},
   "source": [
    "## Instantiate Model using Optimum Intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Optimum Intel can be used to load optimized models from the [Hugging Face Hub](https://huggingface.co/docs/optimum/intel/hf.co/models) and create pipelines to run an inference with OpenVINO Runtime using Hugging Face APIs. The Optimum Inference models are API compatible with Hugging Face Transformers models.  This means we just need to replace `AutoModelForXxx` class with the corresponding `OVModelForXxx` class.\n",
    "\n",
    "Below is an example of the RedPajama model\n",
    "\n",
    "```diff\n",
    "-from transformers import AutoModelForCausalLM\n",
    "+from optimum.intel.openvino import OVModelForCausalLM\n",
    "from transformers import AutoTokenizer, pipeline\n",
    "\n",
    "model_id = \"togethercomputer/RedPajama-INCITE-Chat-3B-v1\"\n",
    "-model = AutoModelForCausalLM.from_pretrained(model_id)\n",
    "+model = OVModelForCausalLM.from_pretrained(model_id, export=True)\n",
    "```\n",
    "\n",
    "Model class initialization starts with calling `from_pretrained` method. When downloading and converting Transformers model, the parameter `export=True` should be added. We can save the converted model for the next usage with the `save_pretrained` method.\n",
    "Tokenizer class and pipelines API are compatible with Optimum models.\n",
    "\n",
    "To optimize the generation process and use memory more efficiently, the `use_cache=True` option is enabled. Since the output side is auto-regressive, an output token hidden state remains the same once computed for every further generation step. Therefore, recomputing it every time you want to generate a new token seems wasteful. With the cache, the model saves the hidden state once it has been computed. The model only computes the one for the most recently generated output token at each time step, re-using the saved ones for hidden tokens. This reduces the generation complexity from $O(n^3)$ to $O(n^2)$ for a transformer model. More details about how it works can be found in this [article](https://scale.com/blog/pytorch-improvements#Text%20Translation). With this option, the model gets the previous step's hidden states (cached attention keys and values) as input and additionally provides hidden states for the current step as output. It means for all next iterations, it is enough to provide only a new token obtained from the previous step and cached key values to get the next token prediction. \n",
    "\n",
    "In our case, MPT, Qwen and ChatGLM model currently is not covered by Optimum Intel, we will convert it manually and create wrapper compatible with Optimum Intel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8cd910c2",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:27:21.307280Z",
     "iopub.status.busy": "2024-03-14T03:27:21.306715Z",
     "iopub.status.idle": "2024-03-14T03:27:31.569250Z",
     "shell.execute_reply": "2024-03-14T03:27:31.568160Z",
     "shell.execute_reply.started": "2024-03-14T03:27:21.307238Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-14 11:27:23,734 - modelscope - INFO - PyTorch version 2.1.2+cpu Found.\n",
      "2024-03-14 11:27:23,738 - modelscope - INFO - TensorFlow version 2.14.0 Found.\n",
      "2024-03-14 11:27:23,740 - modelscope - INFO - Loading ast index from /mnt/workspace/.cache/modelscope/ast_indexer\n",
      "2024-03-14 11:27:23,741 - modelscope - INFO - No valid ast index found from /mnt/workspace/.cache/modelscope/ast_indexer, generating ast index from prebuilt!\n",
      "2024-03-14 11:27:23,802 - modelscope - INFO - Loading done! Current index file version is 1.12.0, with md5 f2107f07efd63ab20c59bbe69eed6202 and a total number of 964 components indexed\n",
      "2024-03-14 11:27:25.270714: I tensorflow/core/util/port.cc:111] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-03-14 11:27:25.273705: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2024-03-14 11:27:25.313913: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-03-14 11:27:25.313950: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-03-14 11:27:25.313978: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-03-14 11:27:25.323778: I tensorflow/tsl/cuda/cudart_stub.cc:28] Could not find cuda drivers on your machine, GPU will not be used.\n",
      "2024-03-14 11:27:25.325132: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-03-14 11:27:26.491214: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, tensorflow, onnx, openvino\n"
     ]
    }
   ],
   "source": [
    "from modelscope import AutoModelForCausalLM, AutoConfig\n",
    "from optimum.intel.openvino import OVModelForCausalLM\n",
    "import openvino as ov\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "import torch\n",
    "import logging\n",
    "import nncf\n",
    "import gc\n",
    "from converter import converters, register_configs\n",
    "\n",
    "register_configs()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13694bf8-ee7b-4186-a3e0-a8705be9733c",
   "metadata": {},
   "source": [
    "## Compress model weights\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "The Weights Compression algorithm is aimed at compressing the weights of the models and can be used to optimize the model footprint and performance of large models where the size of weights is relatively larger than the size of activations, for example, Large Language Models (LLM). Compared to INT8 compression, INT4 compression improves performance even more, but introduces a minor drop in prediction quality.\n",
    "\n",
    "\n",
    "### Weights Compression using Optimum Intel\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "\n",
    "Optimum Intel supports weight compression via NNCF out of the box. For 8-bit compression we pass `load_in_8bit=True` to `from_pretrained()` method of `OVModelForCausalLM`. For 4 bit compression we provide `quantization_config=OVWeightQuantizationConfig(bits=4, ...)` argument containing number of bits and other compression parameters. We will consider how to do it on RedPajama, LLAMA and Zephyr examples.\n",
    "\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU.\n",
    "\n",
    "### Weights Compression using NNCF\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "You also can perform weights compression for OpenVINO models using NNCF directly. `nncf.compress_weights` function accepts OpenVINO model instance and compresses its weights for Linear and Embedding layers. We will consider this variant based on MPT model.\n",
    "\n",
    "\n",
    ">**Note**: This tutorial involves conversion model for FP16 and INT4/INT8 weights compression scenarios. It may be memory and time-consuming in the first run. You can manually control the compression precision below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "91eb2ccf",
   "metadata": {
    "collapsed": false,
    "execution": {
     "iopub.execute_input": "2024-03-14T03:28:01.022284Z",
     "iopub.status.busy": "2024-03-14T03:28:01.021684Z",
     "iopub.status.idle": "2024-03-14T03:28:01.041713Z",
     "shell.execute_reply": "2024-03-14T03:28:01.040559Z",
     "shell.execute_reply.started": "2024-03-14T03:28:01.022241Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "23d1dfdea1ee406ba34e64bca7b42e6a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Prepare INT4 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d4aae8b81e134c9bbfc08e24781ede00",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Prepare INT8 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "06e9d376e28740e4993247e24c117468",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Prepare FP16 model')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "prepare_int4_model = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description=\"Prepare INT4 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "prepare_int8_model = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"Prepare INT8 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "prepare_fp16_model = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description=\"Prepare FP16 model\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "display(prepare_int4_model)\n",
    "display(prepare_int8_model)\n",
    "display(prepare_fp16_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "130a037a-7d98-4152-81ea-92ffb01da5a2",
   "metadata": {},
   "source": [
    "We can now save floating point and compressed model variants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c4ef9112",
   "metadata": {
    "collapsed": false,
    "execution": {
     "iopub.execute_input": "2024-03-14T03:28:08.289506Z",
     "iopub.status.busy": "2024-03-14T03:28:08.288838Z",
     "iopub.status.idle": "2024-03-14T03:32:35.453522Z",
     "shell.execute_reply": "2024-03-14T03:32:35.452339Z",
     "shell.execute_reply.started": "2024-03-14T03:28:08.289442Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading: 100%|██████████| 661/661 [00:00<00:00, 1.73MB/s]\n",
      "Downloading: 100%|██████████| 51.0/51.0 [00:00<00:00, 181kB/s]\n",
      "Downloading: 100%|██████████| 206/206 [00:00<00:00, 466kB/s]\n",
      "Downloading: 100%|██████████| 7.11k/7.11k [00:00<00:00, 11.0MB/s]\n",
      "Downloading: 100%|██████████| 1.59M/1.59M [00:00<00:00, 63.1MB/s]\n",
      "Downloading: 100%|█████████▉| 1.15G/1.15G [00:05<00:00, 247MB/s]\n",
      "Downloading: 100%|██████████| 4.15k/4.15k [00:00<00:00, 5.97MB/s]\n",
      "Downloading: 100%|██████████| 6.70M/6.70M [00:00<00:00, 64.3MB/s]\n",
      "Downloading: 100%|██████████| 1.37k/1.37k [00:00<00:00, 3.90MB/s]\n",
      "Downloading: 100%|██████████| 2.65M/2.65M [00:00<00:00, 85.2MB/s]\n",
      "2024-03-14 11:28:21,146 - modelscope - WARNING - Download interval is too small, use local cache\n",
      "This architecture : qwen2 was not validated, only :llama, gpt-bigcode, blenderbot, codegen, bart, opt, gpt2, gpt-neox, marian, gpt-neo, pegasus, bloom, blenderbot-small architectures were validated, use at your own risk.\n",
      "Framework not specified. Using pt to export the model.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Using the export variant default. Available variants are:\n",
      "    - default: The default ONNX variant.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Using framework PyTorch: 2.1.2+cpu\n",
      "Overriding 1 configuration item(s)\n",
      "\t- use_cache -> True\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/modeling_utils.py:4193: FutureWarning: `_is_quantized_training_enabled` is going to be deprecated in transformers 4.39.0. Please use `model.hf_quantizer.is_trainable` instead\n",
      "  warnings.warn(\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/modeling_attn_mask_utils.py:114: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if (input_shape[-1] > 1 or self.sliding_window is not None) and self.is_causal:\n",
      "/opt/conda/lib/python3.10/site-packages/optimum/exporters/onnx/model_patcher.py:301: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if past_key_values_length > 0:\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py:126: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if seq_len > self.max_seq_len_cached:\n",
      "/opt/conda/lib/python3.10/site-packages/transformers/models/qwen2/modeling_qwen2.py:686: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
      "  if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0a10c7b12a68465cb8847f6a952a3b77",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "869afd684c604f028f3dd0ebebb71c92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Configuration saved in qwen1.5-0.5b-chat/INT4_compressed_weights/openvino_config.json\n"
     ]
    }
   ],
   "source": [
    "from optimum.intel import OVWeightQuantizationConfig\n",
    "from modelscope import snapshot_download\n",
    "\n",
    "nncf.set_log_level(logging.ERROR)\n",
    "\n",
    "pt_model_id = model_configuration[\"model_id\"]\n",
    "pt_model_name = model_id.value.split(\"-\")[0]\n",
    "model_dir = snapshot_download(pt_model_id)\n",
    "model_type = AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True).model_type\n",
    "fp16_model_dir = Path(model_id.value) / \"FP16\"\n",
    "int8_model_dir = Path(model_id.value) / \"INT8_compressed_weights\"\n",
    "int4_model_dir = Path(model_id.value) / \"INT4_compressed_weights\"\n",
    "\n",
    "\n",
    "def convert_to_fp16():\n",
    "    if (fp16_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    if not model_configuration[\"remote\"]:\n",
    "        remote_code = model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\": True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            model_dir, export=True, compile=False, load_in_8bit=False, **model_kwargs\n",
    "        )\n",
    "        ov_model.half()\n",
    "        ov_model.save_pretrained(fp16_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        model_kwargs = {}\n",
    "        if \"revision\" in model_configuration:\n",
    "            model_kwargs[\"revision\"] = model_configuration[\"revision\"]\n",
    "        model = AutoModelForCausalLM.from_pretrained(\n",
    "            pt_model_id,\n",
    "            torch_dtype=torch.float32,\n",
    "            trust_remote_code=True,\n",
    "            **model_kwargs\n",
    "        )\n",
    "        converters[pt_model_name](model, fp16_model_dir)\n",
    "        del model\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "def convert_to_int8():\n",
    "    if (int8_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    int8_model_dir.mkdir(parents=True, exist_ok=True)\n",
    "    if not model_configuration[\"remote\"]:\n",
    "        remote_code = model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\": True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            model_dir, export=True, compile=False, load_in_8bit=True, **model_kwargs\n",
    "        )\n",
    "        ov_model.save_pretrained(int8_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        convert_to_fp16()\n",
    "        ov_model = ov.Core().read_model(fp16_model_dir / \"openvino_model.xml\")\n",
    "        shutil.copy(fp16_model_dir / \"config.json\", int8_model_dir / \"config.json\")\n",
    "        configuration_file = fp16_model_dir / f\"configuration_{model_type}.py\"\n",
    "        if configuration_file.exists():\n",
    "            shutil.copy(\n",
    "                configuration_file, int8_model_dir / f\"configuration_{model_type}.py\"\n",
    "            )\n",
    "        compressed_model = nncf.compress_weights(ov_model)\n",
    "        ov.save_model(compressed_model, int8_model_dir / \"openvino_model.xml\")\n",
    "        del ov_model\n",
    "        del compressed_model\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "def convert_to_int4():\n",
    "    compression_configs = {\n",
    "        \"zephyr-7b-beta\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"mistral-7b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"minicpm-2b-dpo\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"gemma-2b-it\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"notus-7b-v1\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"neural-chat-7b-v3-1\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 64,\n",
    "            \"ratio\": 0.6,\n",
    "        },\n",
    "        \"llama-2-chat-7b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "        \"gemma-7b-it\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "        \"chatglm2-6b\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.72,\n",
    "        },\n",
    "        \"qwen-7b-chat\": {\n",
    "            \"sym\": True,\n",
    "            \"group_size\": 128, \n",
    "            \"ratio\": 0.6\n",
    "        },\n",
    "        'red-pajama-3b-chat': {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.5,\n",
    "        },\n",
    "        \"default\": {\n",
    "            \"sym\": False,\n",
    "            \"group_size\": 128,\n",
    "            \"ratio\": 0.8,\n",
    "        },\n",
    "    }\n",
    "\n",
    "    model_compression_params = compression_configs.get(\n",
    "        model_id.value, compression_configs[\"default\"]\n",
    "    )\n",
    "    if (int4_model_dir / \"openvino_model.xml\").exists():\n",
    "        return\n",
    "    int4_model_dir.mkdir(parents=True, exist_ok=True)\n",
    "    if not model_configuration[\"remote\"]:\n",
    "        remote_code = model_configuration.get(\"remote_code\", False)\n",
    "        model_kwargs = {}\n",
    "        if remote_code:\n",
    "            model_kwargs = {\n",
    "                \"trust_remote_code\" : True,\n",
    "                \"config\": AutoConfig.from_pretrained(pt_model_id, trust_remote_code=True)\n",
    "            }\n",
    "        ov_model = OVModelForCausalLM.from_pretrained(\n",
    "            model_dir, export=True, compile=False,\n",
    "            quantization_config=OVWeightQuantizationConfig(bits=4, **model_compression_params),\n",
    "            **model_kwargs\n",
    "        )\n",
    "        ov_model.save_pretrained(int4_model_dir)\n",
    "        del ov_model\n",
    "    else:\n",
    "        convert_to_fp16()\n",
    "        ov_model = ov.Core().read_model(fp16_model_dir / \"openvino_model.xml\")\n",
    "        shutil.copy(fp16_model_dir / \"config.json\", int4_model_dir / \"config.json\")\n",
    "        configuration_file = fp16_model_dir / f\"configuration_{model_type}.py\"\n",
    "        if configuration_file.exists():\n",
    "            shutil.copy(\n",
    "                configuration_file, int4_model_dir / f\"configuration_{model_type}.py\"\n",
    "            )\n",
    "        mode = nncf.CompressWeightsMode.INT4_SYM if model_compression_params[\"sym\"] else \\\n",
    "            nncf.CompressWeightsMode.INT4_ASYM\n",
    "        del model_compression_params[\"sym\"]\n",
    "        compressed_model = nncf.compress_weights(ov_model, mode=mode, **model_compression_params)\n",
    "        ov.save_model(compressed_model, int4_model_dir / \"openvino_model.xml\")\n",
    "        del ov_model\n",
    "        del compressed_model\n",
    "    gc.collect()\n",
    "\n",
    "\n",
    "if prepare_fp16_model.value:\n",
    "    convert_to_fp16()\n",
    "if prepare_int8_model.value:\n",
    "    convert_to_int8()\n",
    "if prepare_int4_model.value:\n",
    "    convert_to_int4()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "671a17d4",
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Let's compare model size for different compression types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "281f1d07-998e-4e13-ba95-0264564ede82",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:32:45.709230Z",
     "iopub.status.busy": "2024-03-14T03:32:45.708608Z",
     "iopub.status.idle": "2024-03-14T03:32:45.717951Z",
     "shell.execute_reply": "2024-03-14T03:32:45.716733Z",
     "shell.execute_reply.started": "2024-03-14T03:32:45.709182Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Size of model with INT4 compressed weights is 346.55 MB\n"
     ]
    }
   ],
   "source": [
    "fp16_weights = fp16_model_dir / \"openvino_model.bin\"\n",
    "int8_weights = int8_model_dir / \"openvino_model.bin\"\n",
    "int4_weights = int4_model_dir / \"openvino_model.bin\"\n",
    "\n",
    "if fp16_weights.exists():\n",
    "    print(f\"Size of FP16 model is {fp16_weights.stat().st_size / 1024 / 1024:.2f} MB\")\n",
    "for precision, compressed_weights in zip([8, 4], [int8_weights, int4_weights]):\n",
    "    if compressed_weights.exists():\n",
    "        print(\n",
    "            f\"Size of model with INT{precision} compressed weights is {compressed_weights.stat().st_size / 1024 / 1024:.2f} MB\"\n",
    "        )\n",
    "    if compressed_weights.exists() and fp16_weights.exists():\n",
    "        print(\n",
    "            f\"Compression rate for INT{precision} model: {fp16_weights.stat().st_size / compressed_weights.stat().st_size:.3f}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d62f9f4-5434-4550-b372-c86b5a5089d5",
   "metadata": {},
   "source": [
    "## Select device for inference and model variant\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    ">**Note**: There may be no speedup for INT4/INT8 compressed models on dGPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "837b4a3b-ccc3-4004-9577-2b2c7b802dea",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:32:49.523234Z",
     "iopub.status.busy": "2024-03-14T03:32:49.522601Z",
     "iopub.status.idle": "2024-03-14T03:32:49.576037Z",
     "shell.execute_reply": "2024-03-14T03:32:49.574867Z",
     "shell.execute_reply.started": "2024-03-14T03:32:49.523176Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "55a4a98e4a48497b8adee1066d79236e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Device:', options=('CPU', 'AUTO'), value='CPU')"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "core = ov.Core()\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value=\"CPU\",\n",
    "    description=\"Device:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd55ade7-0445-47e1-90f0-72b82da016ca",
   "metadata": {},
   "source": [
    "The cell below create `OVMPTModel` and `OVCHATGLM2Model` wrapper based on `OVModelForCausalLM` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5333ab9b-ff5d-4a7f-bcdc-9cca5d56dc0a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:32:54.415576Z",
     "iopub.status.busy": "2024-03-14T03:32:54.414934Z",
     "iopub.status.idle": "2024-03-14T03:32:54.427292Z",
     "shell.execute_reply": "2024-03-14T03:32:54.425920Z",
     "shell.execute_reply.started": "2024-03-14T03:32:54.415533Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ov_llm_model import model_classes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c53001e7-615f-4eb5-b831-4e2b2ff32826",
   "metadata": {
    "tags": []
   },
   "source": [
    "The cell below demonstrates how to instantiate model based on selected variant of model weights and inference device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3536a1a7",
   "metadata": {
    "collapsed": false,
    "execution": {
     "iopub.execute_input": "2024-03-14T03:32:57.614418Z",
     "iopub.status.busy": "2024-03-14T03:32:57.613797Z",
     "iopub.status.idle": "2024-03-14T03:32:57.627164Z",
     "shell.execute_reply": "2024-03-14T03:32:57.626020Z",
     "shell.execute_reply.started": "2024-03-14T03:32:57.614370Z"
    },
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "80151f1779fa4bb4b1b7fa7f76bde991",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Model to run:', options=('INT4',), value='INT4')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "available_models = []\n",
    "if int4_model_dir.exists():\n",
    "    available_models.append(\"INT4\")\n",
    "if int8_model_dir.exists():\n",
    "    available_models.append(\"INT8\")\n",
    "if fp16_model_dir.exists():\n",
    "    available_models.append(\"FP16\")\n",
    "\n",
    "model_to_run = widgets.Dropdown(\n",
    "    options=available_models,\n",
    "    value=available_models[0],\n",
    "    description=\"Model to run:\",\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "model_to_run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7a041101-7336-40fd-96c9-cd298015a0f3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:33:02.652512Z",
     "iopub.status.busy": "2024-03-14T03:33:02.651894Z",
     "iopub.status.idle": "2024-03-14T03:33:09.539929Z",
     "shell.execute_reply": "2024-03-14T03:33:09.538632Z",
     "shell.execute_reply.started": "2024-03-14T03:33:02.652445Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model from qwen1.5-0.5b-chat/INT4_compressed_weights\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "The argument `trust_remote_code` is to be used along with export=True. It will be ignored.\n",
      "Compiling the model to CPU ...\n"
     ]
    }
   ],
   "source": [
    "from modelscope import AutoTokenizer\n",
    "\n",
    "if model_to_run.value == \"INT4\":\n",
    "    model_dir = int4_model_dir\n",
    "elif model_to_run.value == \"INT8\":\n",
    "    model_dir = int8_model_dir\n",
    "else:\n",
    "    model_dir = fp16_model_dir\n",
    "print(f\"Loading model from {model_dir}\")\n",
    "\n",
    "\n",
    "ov_config = {\"PERFORMANCE_HINT\": \"LATENCY\", \"NUM_STREAMS\": \"1\", \"CACHE_DIR\": \"\"}\n",
    "\n",
    "# On a GPU device a model is executed in FP16 precision. For red-pajama-3b-chat model there known accuracy\n",
    "# issues caused by this, which we avoid by setting precision hint to \"f32\".\n",
    "if model_id.value == \"red-pajama-3b-chat\" and \"GPU\" in core.available_devices and device.value in [\"GPU\", \"AUTO\"]:\n",
    "    ov_config[\"INFERENCE_PRECISION_HINT\"] = \"f32\"\n",
    "\n",
    "model_name = model_configuration[\"model_id\"]\n",
    "class_key = model_id.value.split(\"-\")[0]\n",
    "tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
    "\n",
    "model_class = (\n",
    "    OVModelForCausalLM\n",
    "    if not model_configuration[\"remote\"]\n",
    "    else model_classes[class_key]\n",
    ")\n",
    "ov_model = model_class.from_pretrained(\n",
    "    model_dir,\n",
    "    device=device.value,\n",
    "    ov_config=ov_config,\n",
    "    config=AutoConfig.from_pretrained(model_dir, trust_remote_code=True),\n",
    "    trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8f6f7596-5677-4931-875b-aaabfa23cabc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:33:09.838436Z",
     "iopub.status.busy": "2024-03-14T03:33:09.837846Z",
     "iopub.status.idle": "2024-03-14T03:33:12.858253Z",
     "shell.execute_reply": "2024-03-14T03:33:12.857084Z",
     "shell.execute_reply.started": "2024-03-14T03:33:09.838389Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:151645 for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2 + 2 = 3\n"
     ]
    }
   ],
   "source": [
    "tokenizer_kwargs = model_configuration.get(\"tokenizer_kwargs\", {})\n",
    "test_string = \"2 + 2 =\"\n",
    "input_tokens = tok(test_string, return_tensors=\"pt\", **tokenizer_kwargs)\n",
    "answer = ov_model.generate(**input_tokens, max_new_tokens=2)\n",
    "print(tok.batch_decode(answer, skip_special_tokens=True)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d622d0-be46-47c0-a762-88cb50ab15a9",
   "metadata": {},
   "source": [
    "## Run Chatbot\n",
    "[back to top ⬆️](#Table-of-contents:)\n",
    "\n",
    "Now, when model created, we can setup Chatbot interface using [Gradio](https://www.gradio.app/).\n",
    "The diagram below illustrates how the chatbot pipeline works\n",
    "\n",
    "![generation pipeline](https://user-images.githubusercontent.com/29454499/255523209-d9336491-c7ba-4dc1-98f0-07f23743ce89.png)\n",
    "\n",
    "As can be seen, the pipeline very similar to instruction-following with only changes that previous conversation history additionally passed as input with next user question for getting wider input context. On the first iteration, the user provided instructions joined to conversation history (if exists) converted to token ids using a tokenizer, then prepared input provided to the model. The model generates probabilities for all tokens in logits format  The way the next token will be selected over predicted probabilities is driven by the selected decoding methodology. You can find more information about the most popular decoding methods in this [blog](https://huggingface.co/blog/how-to-generate). The result generation updates conversation history for next conversation step. it makes stronger connection of next question with previously provided and allows user to make clarifications regarding previously provided answers."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "725544ea-05ec-40d7-bbbc-1dc87cf57d04",
   "metadata": {},
   "source": [
    "There are several parameters that can control text generation quality: \n",
    "  * `Temperature` is a parameter used to control the level of creativity in AI-generated text. By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.  \n",
    "  Consider the following example: The AI model has to complete the sentence \"The cat is ____.\" with the following token probabilities:  \n",
    "\n",
    "    playing: 0.5  \n",
    "    sleeping: 0.25  \n",
    "    eating: 0.15  \n",
    "    driving: 0.05  \n",
    "    flying: 0.05  \n",
    "\n",
    "    - **Low temperature** (e.g., 0.2): The AI model becomes more focused and deterministic, choosing tokens with the highest probability, such as \"playing.\"  \n",
    "    - **Medium temperature** (e.g., 1.0): The AI model maintains a balance between creativity and focus, selecting tokens based on their probabilities without significant bias, such as \"playing,\" \"sleeping,\" or \"eating.\"  \n",
    "    - **High temperature** (e.g., 2.0): The AI model becomes more adventurous, increasing the chances of selecting less likely tokens, such as \"driving\" and \"flying.\"\n",
    "  * `Top-p`, also known as nucleus sampling, is a parameter used to control the range of tokens considered by the AI model based on their cumulative probability. By adjusting the `top-p` value, you can influence the AI model's token selection, making it more focused or diverse.\n",
    "  Using the same example with the cat, consider the following top_p settings:  \n",
    "    - **Low top_p** (e.g., 0.5): The AI model considers only tokens with the highest cumulative probability, such as \"playing.\"  \n",
    "    - **Medium top_p** (e.g., 0.8): The AI model considers tokens with a higher cumulative probability, such as \"playing,\" \"sleeping,\" and \"eating.\"  \n",
    "    - **High top_p** (e.g., 1.0): The AI model considers all tokens, including those with lower probabilities, such as \"driving\" and \"flying.\" \n",
    "  * `Top-k` is an another popular sampling strategy. In comparison with Top-P, which chooses from the smallest possible set of words whose cumulative probability exceeds the probability P, in Top-K sampling K most likely next words are filtered and the probability mass is redistributed among only those K next words. In our example with cat, if k=3, then only \"playing\", \"sleeping\" and \"eating\" will be taken into account as possible next word.\n",
    "  * `Repetition Penalty` This parameter can help penalize tokens based on how frequently they occur in the text, including the input prompt. A token that has already appeared five times is penalized more heavily than a token that has appeared only one time. A value of 1 means that there is no penalty and values larger than 1 discourage repeated tokens."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "01f8f7f8-072e-45dc-b7c9-18d8c3c47754",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-03-14T03:40:50.822317Z",
     "iopub.status.busy": "2024-03-14T03:40:50.821354Z",
     "iopub.status.idle": "2024-03-14T03:40:53.192985Z",
     "shell.execute_reply": "2024-03-14T03:40:53.191764Z",
     "shell.execute_reply.started": "2024-03-14T03:40:50.822265Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7865\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7865/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:151645 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "from threading import Event, Thread\n",
    "from uuid import uuid4\n",
    "from typing import List, Tuple\n",
    "import gradio as gr\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    StoppingCriteria,\n",
    "    StoppingCriteriaList,\n",
    "    TextIteratorStreamer,\n",
    ")\n",
    "import os\n",
    "name = os.environ['JUPYTER_NAME']\n",
    "region = os.environ[\"dsw_region\"]\n",
    "\n",
    "host = \"dsw-gateway-{region}.data.aliyun.com\".format(region=region)\n",
    "host\n",
    "\n",
    "\n",
    "model_name = model_configuration[\"model_id\"]\n",
    "start_message = model_configuration[\"start_message\"]\n",
    "history_template = model_configuration.get(\"history_template\")\n",
    "current_message_template = model_configuration.get(\"current_message_template\")\n",
    "stop_tokens = model_configuration.get(\"stop_tokens\")\n",
    "tokenizer_kwargs = model_configuration.get(\"tokenizer_kwargs\", {})\n",
    "\n",
    "chinese_examples = [\n",
    "    [\"你好!\"],\n",
    "    [\"你是谁?\"],\n",
    "    [\"请介绍一下上海\"],\n",
    "    [\"请介绍一下英特尔公司\"],\n",
    "    [\"晚上睡不着怎么办？\"],\n",
    "    [\"给我讲一个年轻人奋斗创业最终取得成功的故事。\"],\n",
    "    [\"给这个故事起一个标题。\"],\n",
    "]\n",
    "\n",
    "english_examples = [\n",
    "    [\"Hello there! How are you doing?\"],\n",
    "    [\"What is OpenVINO?\"],\n",
    "    [\"Who are you?\"],\n",
    "    [\"Can you explain to me briefly what is Python programming language?\"],\n",
    "    [\"Explain the plot of Cinderella in a sentence.\"],\n",
    "    [\"What are some common mistakes to avoid when writing code?\"],\n",
    "    [\n",
    "        \"Write a 100-word blog post on “Benefits of Artificial Intelligence and OpenVINO“\"\n",
    "    ],\n",
    "]\n",
    "\n",
    "japanese_examples = [\n",
    "    [\"こんにちは！調子はどうですか?\"],\n",
    "    [\"OpenVINOとは何ですか?\"],\n",
    "    [\"あなたは誰ですか?\"],\n",
    "    [\"Pythonプログラミング言語とは何か簡単に説明してもらえますか?\"],\n",
    "    [\"シンデレラのあらすじを一文で説明してください。\"],\n",
    "    [\"コードを書くときに避けるべきよくある間違いは何ですか?\"],\n",
    "    [\"人工知能と「OpenVINOの利点」について100語程度のブログ記事を書いてください。\"],\n",
    "]\n",
    "\n",
    "examples = (\n",
    "    chinese_examples\n",
    "    if (model_language.value == \"Chinese\")\n",
    "    else japanese_examples\n",
    "    if (model_language.value == \"Japanese\")\n",
    "    else english_examples\n",
    ")\n",
    "\n",
    "max_new_tokens = 256\n",
    "\n",
    "\n",
    "class StopOnTokens(StoppingCriteria):\n",
    "    def __init__(self, token_ids):\n",
    "        self.token_ids = token_ids\n",
    "\n",
    "    def __call__(\n",
    "        self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs\n",
    "    ) -> bool:\n",
    "        for stop_id in self.token_ids:\n",
    "            if input_ids[0][-1] == stop_id:\n",
    "                return True\n",
    "        return False\n",
    "\n",
    "\n",
    "if stop_tokens is not None:\n",
    "    if isinstance(stop_tokens[0], str):\n",
    "        stop_tokens = tok.convert_tokens_to_ids(stop_tokens)\n",
    "\n",
    "    stop_tokens = [StopOnTokens(stop_tokens)]\n",
    "\n",
    "\n",
    "def default_partial_text_processor(partial_text: str, new_text: str):\n",
    "    \"\"\"\n",
    "    helper for updating partially generated answer, used by default\n",
    "\n",
    "    Params:\n",
    "      partial_text: text buffer for storing previosly generated text\n",
    "      new_text: text update for the current step\n",
    "    Returns:\n",
    "      updated text string\n",
    "\n",
    "    \"\"\"\n",
    "    partial_text += new_text\n",
    "    return partial_text\n",
    "\n",
    "\n",
    "text_processor = model_configuration.get(\n",
    "    \"partial_text_processor\", default_partial_text_processor\n",
    ")\n",
    "\n",
    "\n",
    "def convert_history_to_token(history: List[Tuple[str, str]]):\n",
    "    \"\"\"\n",
    "    function for conversion history stored as list pairs of user and assistant messages to tokens according to model expected conversation template\n",
    "    Params:\n",
    "      history: dialogue history\n",
    "    Returns:\n",
    "      history in token format\n",
    "    \"\"\"\n",
    "    if pt_model_name == \"baichuan2\":\n",
    "        system_tokens = tok.encode(start_message)\n",
    "        history_tokens = []\n",
    "        for (old_query, response) in history[:-1]:\n",
    "            round_tokens = []\n",
    "            round_tokens.append(195)\n",
    "            round_tokens.extend(tok.encode(old_query))\n",
    "            round_tokens.append(196)\n",
    "            round_tokens.extend(tok.encode(response))\n",
    "            history_tokens = round_tokens + history_tokens\n",
    "        input_tokens = system_tokens + history_tokens\n",
    "        input_tokens.append(195)\n",
    "        input_tokens.extend(tok.encode(history[-1][0]))\n",
    "        input_tokens.append(196)\n",
    "        input_token = torch.LongTensor([input_tokens])\n",
    "    elif history_template is None:\n",
    "        messages = [{\"role\": \"system\", \"content\": start_message}]\n",
    "        for idx, (user_msg, model_msg) in enumerate(history):\n",
    "            if idx == len(history) - 1 and not model_msg:\n",
    "                messages.append({\"role\": \"user\", \"content\": user_msg})\n",
    "                break\n",
    "            if user_msg:\n",
    "                messages.append({\"role\": \"user\", \"content\": user_msg})\n",
    "            if model_msg:\n",
    "                messages.append({\"role\": \"assistant\", \"content\": model_msg})\n",
    "                \n",
    "        input_token = tok.apply_chat_template(messages,\n",
    "                                              add_generation_prompt=True,\n",
    "                                              tokenize=True,\n",
    "                                              return_tensors=\"pt\")\n",
    "    else:\n",
    "        text = start_message + \"\".join(\n",
    "            [\n",
    "                \"\".join(\n",
    "                    [\n",
    "                        history_template.format(\n",
    "                            num=round, user=item[0], assistant=item[1]\n",
    "                        )\n",
    "                    ]\n",
    "                )\n",
    "                for round, item in enumerate(history[:-1])\n",
    "            ]\n",
    "        )\n",
    "        text += \"\".join(\n",
    "            [\n",
    "                \"\".join(\n",
    "                    [\n",
    "                        current_message_template.format(\n",
    "                            num=len(history) + 1,\n",
    "                            user=history[-1][0],\n",
    "                            assistant=history[-1][1],\n",
    "                        )\n",
    "                    ]\n",
    "                )\n",
    "            ]\n",
    "        )\n",
    "        input_token = tok(text, return_tensors=\"pt\", **tokenizer_kwargs).input_ids\n",
    "    return input_token\n",
    "\n",
    "\n",
    "def user(message, history):\n",
    "    \"\"\"\n",
    "    callback function for updating user messages in interface on submit button click\n",
    "\n",
    "    Params:\n",
    "      message: current message\n",
    "      history: conversation history\n",
    "    Returns:\n",
    "      None\n",
    "    \"\"\"\n",
    "    # Append the user's message to the conversation history\n",
    "    return \"\", history + [[message, \"\"]]\n",
    "\n",
    "\n",
    "def bot(history, temperature, top_p, top_k, repetition_penalty, conversation_id):\n",
    "    \"\"\"\n",
    "    callback function for running chatbot on submit button click\n",
    "\n",
    "    Params:\n",
    "      history: conversation history\n",
    "      temperature:  parameter for control the level of creativity in AI-generated text.\n",
    "                    By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.\n",
    "      top_p: parameter for control the range of tokens considered by the AI model based on their cumulative probability.\n",
    "      top_k: parameter for control the range of tokens considered by the AI model based on their cumulative probability, selecting number of tokens with highest probability.\n",
    "      repetition_penalty: parameter for penalizing tokens based on how frequently they occur in the text.\n",
    "      conversation_id: unique conversation identifier.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    # Construct the input message string for the model by concatenating the current system message and conversation history\n",
    "    # Tokenize the messages string\n",
    "    input_ids = convert_history_to_token(history)\n",
    "    if input_ids.shape[1] > 2000:\n",
    "        history = [history[-1]]\n",
    "        input_ids = convert_history_to_token(history)\n",
    "    streamer = TextIteratorStreamer(\n",
    "        tok, timeout=30.0, skip_prompt=True, skip_special_tokens=True\n",
    "    )\n",
    "    generate_kwargs = dict(\n",
    "        input_ids=input_ids,\n",
    "        max_new_tokens=max_new_tokens,\n",
    "        temperature=temperature,\n",
    "        do_sample=temperature > 0.0,\n",
    "        top_p=top_p,\n",
    "        top_k=top_k,\n",
    "        repetition_penalty=repetition_penalty,\n",
    "        streamer=streamer,\n",
    "    )\n",
    "    if stop_tokens is not None:\n",
    "        generate_kwargs[\"stopping_criteria\"] = StoppingCriteriaList(\n",
    "            stop_tokens)\n",
    "\n",
    "    stream_complete = Event()\n",
    "\n",
    "    def generate_and_signal_complete():\n",
    "        \"\"\"\n",
    "        genration function for single thread\n",
    "        \"\"\"\n",
    "        global start_time\n",
    "        ov_model.generate(**generate_kwargs)\n",
    "        stream_complete.set()\n",
    "\n",
    "    t1 = Thread(target=generate_and_signal_complete)\n",
    "    t1.start()\n",
    "\n",
    "    # Initialize an empty string to store the generated text\n",
    "    partial_text = \"\"\n",
    "    for new_text in streamer:\n",
    "        partial_text = text_processor(partial_text, new_text)\n",
    "        history[-1][1] = partial_text\n",
    "        yield history\n",
    "\n",
    "\n",
    "def request_cancel():\n",
    "    ov_model.request.cancel()\n",
    "\n",
    "\n",
    "def get_uuid():\n",
    "    \"\"\"\n",
    "    universal unique identifier for thread\n",
    "    \"\"\"\n",
    "    return str(uuid4())\n",
    "\n",
    "\n",
    "with gr.Blocks(\n",
    "    theme=gr.themes.Soft(),\n",
    "    css=\".disclaimer {font-variant-caps: all-small-caps;}\",\n",
    ") as demo:\n",
    "    conversation_id = gr.State(get_uuid)\n",
    "    gr.Markdown(\n",
    "        f\"\"\"<h1><center>OpenVINO {model_id.value} Chatbot</center></h1>\"\"\")\n",
    "    chatbot = gr.Chatbot(height=500)\n",
    "    with gr.Row():\n",
    "        with gr.Column():\n",
    "            msg = gr.Textbox(\n",
    "                label=\"Chat Message Box\",\n",
    "                placeholder=\"Chat Message Box\",\n",
    "                show_label=False,\n",
    "                container=False,\n",
    "            )\n",
    "        with gr.Column():\n",
    "            with gr.Row():\n",
    "                submit = gr.Button(\"Submit\")\n",
    "                stop = gr.Button(\"Stop\")\n",
    "                clear = gr.Button(\"Clear\")\n",
    "    with gr.Row():\n",
    "        with gr.Accordion(\"Advanced Options:\", open=False):\n",
    "            with gr.Row():\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        temperature = gr.Slider(\n",
    "                            label=\"Temperature\",\n",
    "                            value=0.1,\n",
    "                            minimum=0.0,\n",
    "                            maximum=1.0,\n",
    "                            step=0.1,\n",
    "                            interactive=True,\n",
    "                            info=\"Higher values produce more diverse outputs\",\n",
    "                        )\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        top_p = gr.Slider(\n",
    "                            label=\"Top-p (nucleus sampling)\",\n",
    "                            value=1.0,\n",
    "                            minimum=0.0,\n",
    "                            maximum=1,\n",
    "                            step=0.01,\n",
    "                            interactive=True,\n",
    "                            info=(\n",
    "                                \"Sample from the smallest possible set of tokens whose cumulative probability \"\n",
    "                                \"exceeds top_p. Set to 1 to disable and sample from all tokens.\"\n",
    "                            ),\n",
    "                        )\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        top_k = gr.Slider(\n",
    "                            label=\"Top-k\",\n",
    "                            value=50,\n",
    "                            minimum=0.0,\n",
    "                            maximum=200,\n",
    "                            step=1,\n",
    "                            interactive=True,\n",
    "                            info=\"Sample from a shortlist of top-k tokens — 0 to disable and sample from all tokens.\",\n",
    "                        )\n",
    "                with gr.Column():\n",
    "                    with gr.Row():\n",
    "                        repetition_penalty = gr.Slider(\n",
    "                            label=\"Repetition Penalty\",\n",
    "                            value=1.1,\n",
    "                            minimum=1.0,\n",
    "                            maximum=2.0,\n",
    "                            step=0.1,\n",
    "                            interactive=True,\n",
    "                            info=\"Penalize repetition — 1.0 to disable.\",\n",
    "                        )\n",
    "    gr.Examples(\n",
    "        examples, inputs=msg, label=\"Click on any example and press the 'Submit' button\"\n",
    "    )\n",
    "\n",
    "    submit_event = msg.submit(\n",
    "        fn=user,\n",
    "        inputs=[msg, chatbot],\n",
    "        outputs=[msg, chatbot],\n",
    "        queue=False,\n",
    "    ).then(\n",
    "        fn=bot,\n",
    "        inputs=[\n",
    "            chatbot,\n",
    "            temperature,\n",
    "            top_p,\n",
    "            top_k,\n",
    "            repetition_penalty,\n",
    "            conversation_id,\n",
    "        ],\n",
    "        outputs=chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    submit_click_event = submit.click(\n",
    "        fn=user,\n",
    "        inputs=[msg, chatbot],\n",
    "        outputs=[msg, chatbot],\n",
    "        queue=False,\n",
    "    ).then(\n",
    "        fn=bot,\n",
    "        inputs=[\n",
    "            chatbot,\n",
    "            temperature,\n",
    "            top_p,\n",
    "            top_k,\n",
    "            repetition_penalty,\n",
    "            conversation_id,\n",
    "        ],\n",
    "        outputs=chatbot,\n",
    "        queue=True,\n",
    "    )\n",
    "    stop.click(\n",
    "        fn=request_cancel,\n",
    "        inputs=None,\n",
    "        outputs=None,\n",
    "        cancels=[submit_event, submit_click_event],\n",
    "        queue=False,\n",
    "    )\n",
    "    clear.click(lambda: None, None, chatbot, queue=False)\n",
    "\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "#  demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# if you have any issue to launch on your platform, you can pass share=True to launch method:\n",
    "# demo.launch(share=True)\n",
    "# it creates a publicly shareable link for the interface. Read more in the docs: https://gradio.app/docs/\n",
    "port = 7865\n",
    "root_path = f'/{name}/proxy/{port}'\n",
    "demo.launch(root_path=root_path,server_port=port)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "7b837f9e-4152-4a5c-880a-ed874aa64a74",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Closing server running on port: 5467\n"
     ]
    }
   ],
   "source": [
    "# please uncomment and run this cell for stopping gradio interface\n",
    "# demo.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "openvino_notebooks": {
   "imageUrl": "https://user-images.githubusercontent.com/29454499/255799218-611e7189-8979-4ef5-8a80-5a75e0136b50.png",
   "tags": {
    "categories": [
     "Model Demos",
     "AI Trends"
    ],
    "libraries": [],
    "other": [],
    "tasks": [
     "Text Generation",
     "Conversational"
    ]
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
